from Linac import *
from mpi4py import MPI
import storage
from progressbar import Bar, Percentage, RotatingMarker,ETA, ProgressBar

LinacDesign = 'optimusPlus/'
fileDir = '/home/ainswort/HOMDynamics/LatticeFiles/'+LinacDesign
saveDir = '/data/ainswort/Output/'+LinacDesign+'HOMs/Run1/'
fileName = 'QIbScan-som.h5'
InjFile = '/home/ainswort/HOMDynamics/RunScripts/dtlErrors/'+LinacDesign+'InputData.dat'


comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

numBunches = int(1e6)
tb = 1/352.2e6
QexScan = [i*j for i in np.linspace(1,9,9) for j in [1e5,1e6,1e7]] 
QexScan.sort()
IbScan = np.array([62.5e-3,75e-3]) 
N_tot = 100

fsigS = 10e3
fsigE = 1e3

Cavities = CavitySection(file=fileDir+'TW_Sync.txt',E_in=89.6e6)
Cavities.cavPeriod = np.concatenate([np.ones(26,dtype='int')*2,np.ones(36,dtype='int')*4,np.ones(84,dtype='int')*4])
Cavities.VoltageCorrect()

if rank==0:

    fileOut = storage.LinacDB(saveDir+fileName,int(numBunches),'w')
    for ib in IbScan:
        for Qex in QexScan:

            input = (ib,Qex)
            groupTitle = 'Current: '+str(int(ib*1000))+'mA and Qex: %1.0E' %Qex
            print groupTitle
            groupName = 'Ib'+str(int(ib*1000))+'mA_Qex%1.0E' %Qex
            groupName = groupName.replace('+','')
            group = fileOut.h5file.createGroup('/', groupName, groupTitle)
            PSOUT = storage.createPhaseSpaceOutDB(numBunches)
            PSOUTtable = fileOut.h5file.createTable("/"+groupName, 'PhaseSpaceExit', PSOUT, expectedrows=N_tot, filters=fileOut.filters)


            pbar  = ProgressBar(widgets=['Linacs Tracked: ', Percentage(), ' ', Bar(marker=RotatingMarker()),' ', ETA()],maxval=N_tot).start()
            N = 0
            process = 1
            while process < size and N<N_tot:
                comm.send(input,dest=process,tag=3)
                N+=1
                process+=1
        
            recieved_processes=0
            while recieved_processes<N_tot:
                L=comm.recv(source=MPI.ANY_SOURCE,tag=1)
                fileOut.Write(L,PSOUTtable)
                process = comm.recv(source=MPI.ANY_SOURCE,tag=2) 
                recieved_processes+=1
                pbar.update(recieved_processes)

                if N<N_tot:
                    comm.send(input,dest=process,tag=3)
                    N+=1
            pbar.finish()
            print '\n'

    for process in range(1,size):
        comm.send(-1,dest=process,tag=3)


else:
    while True:
        Input = comm.recv(source=0,tag=3)
        if Input == -1: break
        
        Ib= Input[0]
        homQex= Input[1]
       
        Ibsig = Ib*0.01
        
        P1 = Pulse(numBunches,Cavities.cavNum, tb, Ib)
        P1.loadInjectionNoise(InjFile)
        P1.AddCurrentError(Ibsig,seed=4612)
        
        H1 = HOM(Cavities.cavNum,Qex=homQex)
        H1.loadHomFREQandRQ('/data/ainswort/SPOKE/'+LinacDesign+'Monopoles-RoverQ.dat',np.array([1]),0,26)
        H1.loadHomFREQandRQ('/data/ainswort/ESS067/'+LinacDesign+'Monopoles-RoverQ.dat',np.array([4]),26,36)
        H1.loadHomFREQandRQ('/data/ainswort/ESS086/'+LinacDesign+'Monopoles-RoverQ.dat',np.array([3]),26+36,84)
        H1.RQCorrect(Cavities.W_i,Cavities.cavPeriod)
        H1.AddFrequencyErrorSundelin(352.21e6,0,0,26)
        H1.AddFrequencyErrorSundelin(704.42e6,0,26,62)
        H1.AddFrequencyErrorSundelin(704.42e6,0,62)
 
        
        Lin = Linac(Cavities,P1,H1)
        Lin.Track()

        comm.send(Lin,dest=0,tag=1)
        comm.send(rank,dest=0,tag=2)

